import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_absolute_error, mean_squared_error
import lightgbm as lgb
from datetime import timedelta

# 1. 数据加载与预处理
# 假设数据文件名为sales_data.csv，包含列：日期,sku,销量,商品季节,上新日期,商品类别
df = pd.read_csv('sales_demo_data.csv')
# data = {
#     '日期': pd.date_range(start='2023-01-01', periods=100, freq='D'),
#     '销量': np.random.randint(100, 500, size=100)  # 随机生成100天的销量数据
# }
# df = pd.DataFrame(data,encoding='GBK')

# 转换日期格式
df['日期'] = pd.to_datetime(df['日期'])
df['上新日期'] = pd.to_datetime(df['上新日期'])

# 按sku和日期排序
df = df.sort_values(['sku', '日期']).reset_index(drop=True)


# 2. 生成目标变量（未来7/14/30天销量总和）
def create_target(group, days):
    return group['销量'].shift(-days).rolling(window=days, min_periods=1).sum()


for days in [7, 14, 30]:
    df[f'target_{days}d'] = df.groupby('sku', group_keys=False)['销量'].transform(
        lambda x: create_target(x, days))

# 删除无法计算目标的记录
df = df.dropna(subset=['target_7d', 'target_14d', 'target_30d'])

# 3. 特征工程
# 基础时间特征
df['year'] = df['日期'].dt.year
df['month'] = df['日期'].dt.month
df['day'] = df['日期'].dt.day
df['dayofweek'] = df['日期'].dt.dayofweek
df['quarter'] = df['日期'].dt.quarter

# 商品生命周期特征
df['days_since_launch'] = (df['日期'] - df['上新日期']).dt.days


# 历史统计特征（按sku分组计算）
def create_history_features(group):
    return group['销量'].rolling(7, min_periods=1).sum().shift(1)


for window in [7, 14, 30]:
    df[f'past_{window}d_sum'] = df.groupby('sku', group_keys=False)['销量'].transform(
        lambda x: x.rolling(window, min_periods=1).sum().shift(1))

# 滞后特征
for lag in [1, 3, 7, 14]:
    df[f'lag_{lag}'] = df.groupby('sku', group_keys=False)['销量'].shift(lag)

# 分类特征处理
cat_features = ['sku', '商品季节', '商品类别']
for col in cat_features:
    df[col] = df[col].astype('category')

# 4. 特征选择
features = ['year', 'month', 'day', 'dayofweek', 'quarter',
            'days_since_launch', 'past_7d_sum', 'past_14d_sum', 'past_30d_sum',
            'lag_1', 'lag_3', 'lag_7', 'lag_14'] + cat_features
targets = ['target_7d', 'target_14d', 'target_30d']

# 5. 数据集划分（按时间序列划分）
split_date = df['日期'].max() - timedelta(days=60)
train = df[df['日期'] <= split_date]
test = df[df['日期'] > split_date]

X_train, X_test = train[features], test[features]
y_train = {t: train[t] for t in targets}
y_test = {t: test[t] for t in targets}

# 6. 模型训练
params = {
    'objective': 'regression',
    'metric': 'mae',
    'num_leaves': 31,
    'learning_rate': 0.05,
    'feature_fraction': 0.9,
    'force_col_wise': True
}

models = {}
for target in targets:
    train_data = lgb.Dataset(X_train, label=y_train[target],
                             categorical_feature=cat_features)
    models[target] = lgb.train(params, train_data, num_boost_round=100)

# 7. 模型评估
for target in targets:
    pred = models[target].predict(X_test)
    print(f'\n评估结果 [{target}]:')
    print(f'MAE: {mean_absolute_error(y_test[target], pred):.2f}')
    print(f'MSE: {mean_squared_error(y_test[target], pred):.2f}')


# 8. 未来预测函数
def predict_future(sku_list, predict_date, days_ahead):
    """
    sku_list: 需要预测的SKU列表
    predict_date: 预测基准日期（datetime格式）
    days_ahead: 需要预测的天数列表，如[7, 14, 30]
    """
    predictions = []
    for sku in sku_list:
        # 获取该SKU最新数据
        sku_data = df[df['sku'] == sku]
        if sku_data.empty:
            continue

        # 生成特征
        latest = sku_data.iloc[-1].copy()
        features = {
            'year': predict_date.year,
            'month': predict_date.month,
            'day': predict_date.day,
            'dayofweek': predict_date.weekday(),
            'quarter': (predict_date.month - 1) // 3 + 1,
            'days_since_launch': (predict_date - latest['上新日期']).days,
            'past_7d_sum': latest['past_7d_sum'],
            'past_14d_sum': latest['past_14d_sum'],
            'past_30d_sum': latest['past_30d_sum'],
            'lag_1': latest['销量'],
            'lag_3': sku_data.iloc[-3]['销量'] if len(sku_data) >= 3 else 0,
            'lag_7': sku_data.iloc[-7]['销量'] if len(sku_data) >= 7 else 0,
            'lag_14': sku_data.iloc[-14]['销量'] if len(sku_data) >= 14 else 0,
            'sku': sku,
            '商品季节': latest['商品季节'],
            '商品类别': latest['商品类别']
        }

        # 转换为DataFrame并处理类别类型
        future_df = pd.DataFrame([features])
        for col in cat_features:
            future_df[col] = future_df[col].astype('category')

        # 进行预测
        preds = {}
        for days in days_ahead:
            model = models.get(f'target_{days}d')
            if model:
                preds[f'{days}d'] = model.predict(future_df)[0]

        predictions.append({
            'sku': sku,
            'predict_date': predict_date,
            **preds
        })

    return pd.DataFrame(predictions)


# 使用示例：预测2023-12-31之后7/14/30天销量
future_date = pd.to_datetime('2023-12-31')
sku_list = df['sku'].unique()[:5]  # 示例预测前5个SKU
predictions = predict_future(sku_list, future_date, [7, 14, 30])

print("\n未来销量预测结果:")
print(predictions)